package icmp

import (
	"fmt"
	"net"
	"time"

	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"
)

const (
	ICMPv4ProtocolNumber = 1  // ICMP for IPv4
	ICMPv6ProtocolNumber = 58 // ICMP for IPv6
)

func SendICMPRequest(address string) error {
	// Determine if the address is IPv4 or IPv6
	ip := net.ParseIP(address)
	if ip == nil {
		return fmt.Errorf("invalid IP address: %s", address)
	}

	var conn *icmp.PacketConn
	var err error
	if ip.To4() != nil {
		// IPv4 connection
		conn, err = icmp.ListenPacket("ip4:icmp", "0.0.0.0")
		if err != nil {
			return fmt.Errorf("failed to listen for ICMPv4: %w", err)
		}
		defer conn.Close()
		return sendICMPv4(conn, ip.String())
	} else {
		// IPv6 connection
		conn, err = icmp.ListenPacket("ip6:ipv6-icmp", "::")
		if err != nil {
			return fmt.Errorf("failed to listen for ICMPv6: %w", err)
		}
		defer conn.Close()
		return sendICMPv6(conn, ip.String())
	}
}

func sendICMPv4(conn *icmp.PacketConn, address string) error {
	var rttTimes []time.Duration
	packetCount := 10
	for i := 0; i < packetCount; i++ {
		msg := icmp.Message{
			Type: ipv4.ICMPTypeEcho, Code: 0,
			Body: &icmp.Echo{
				ID:   1,     // Identifier
				Seq:  i + 1, // Sequence number
				Data: []byte("HELLO-ICMP"),
			},
		}

		bytes, err := msg.Marshal(nil)
		if err != nil {
			return fmt.Errorf("failed to marshal ICMPv4 message: %w", err)
		}

		start := time.Now()
		_, err = conn.WriteTo(bytes, &net.IPAddr{IP: net.ParseIP(address)})
		if err != nil {
			return fmt.Errorf("failed to send ICMPv4 packet: %w", err)
		}

		reply := make([]byte, 1500)
		n, _, err := conn.ReadFrom(reply)
		if err != nil {
			return fmt.Errorf("failed to receive ICMPv4 reply: %w", err)
		}

		parsedMsg, err := icmp.ParseMessage(ICMPv4ProtocolNumber, reply[:n])
		if err != nil {
			return fmt.Errorf("failed to parse ICMPv4 reply: %w", err)
		}

		if parsedMsg.Type == ipv4.ICMPTypeEchoReply {
			elapsed := time.Since(start)
			rttTimes = append(rttTimes, elapsed)
			fmt.Printf("len=28 ip=%s icmp_seq=%d rtt=%.2f ms\n", address, i, float64(elapsed.Microseconds())/1000)
		} else {
			fmt.Printf("received unexpected ICMPv4 message: %+v\n", parsedMsg)
		}
		time.Sleep(1 * time.Second)
	}

	printICMPStatistics(rttTimes)
	return nil
}

func sendICMPv6(conn *icmp.PacketConn, address string) error {
	var rttTimes []time.Duration
	packetCount := 10
	for i := 0; i < packetCount; i++ {
		msg := icmp.Message{
			Type: ipv6.ICMPTypeEchoRequest, Code: 0,
			Body: &icmp.Echo{
				ID:   1,     // Identifier
				Seq:  i + 1, // Sequence number
				Data: []byte("HELLO-ICMP"),
			},
		}

		bytes, err := msg.Marshal(nil)
		if err != nil {
			return fmt.Errorf("failed to marshal ICMPv6 message: %w", err)
		}

		start := time.Now()
		_, err = conn.WriteTo(bytes, &net.IPAddr{IP: net.ParseIP(address)})
		if err != nil {
			return fmt.Errorf("failed to send ICMPv6 packet: %w", err)
		}

		reply := make([]byte, 1500)
		n, _, err := conn.ReadFrom(reply)
		if err != nil {
			return fmt.Errorf("failed to receive ICMPv6 reply: %w", err)
		}

		parsedMsg, err := icmp.ParseMessage(ICMPv6ProtocolNumber, reply[:n])
		if err != nil {
			return fmt.Errorf("failed to parse ICMPv6 reply: %w", err)
		}

		if parsedMsg.Type == ipv6.ICMPTypeEchoReply {
			elapsed := time.Since(start)
			rttTimes = append(rttTimes, elapsed)
			fmt.Printf("len=28 ip=%s icmp_seq=%d rtt=%.2f ms\n", address, i, float64(elapsed.Microseconds())/1000)
		} else {
			fmt.Printf("received unexpected ICMPv6 message: %+v\n", parsedMsg)
		}
		time.Sleep(1 * time.Second)
	}

	printICMPStatistics(rttTimes)
	return nil
}

func printICMPStatistics(rttTimes []time.Duration) {
	packetCount := len(rttTimes)
	if packetCount == 0 {
		fmt.Println("No packets received.")
		return
	}

	var min, max, total time.Duration
	min = rttTimes[0]
	for _, rtt := range rttTimes {
		if rtt < min {
			min = rtt
		}
		if rtt > max {
			max = rtt
		}
		total += rtt
	}

	avg := total / time.Duration(packetCount)
	fmt.Printf("--- ping statistics ---\n")
	fmt.Printf("%d packets transmitted, %d packets received, %.0f%% packet loss\n", packetCount, packetCount, 0.0)
	fmt.Printf("round-trip min/avg/max = %.2f/%.2f/%.2f ms\n", float64(min.Microseconds())/1000, float64(avg.Microseconds())/1000, float64(max.Microseconds())/1000)
}
